import spacy
from tqdm import tqdm

# 加载spaCy中文模型（需要提前下载：python -m spacy download zh_core_web_sm）
nlp = spacy.load("zh_core_web_sm")


def split_text(pdf_contents, chunk_size=200, chunk_overlap=50):
    """
    将PDF提取的文本切分成小块（chunk）

    参数:
        pdf_contents: 从read_pdf函数得到的结果（列表，包含页码和文本）
        chunk_size: 每个片段的大致长度（字符数）
        chunk_overlap: 片段之间的重叠字符数（保持上下文连贯）

    返回:
        列表，每个元素是字典，包含"page_num"（页码）和"chunk"（文本片段）
    """
    chunks = []  # 存储所有切分后的片段

    for item in tqdm(pdf_contents, desc="切分文本中"):
        page_num = item["page_num"]
        full_text = item["text"]

        # 用spaCy处理文本，进行分词和句子分割
        doc = nlp(full_text)
        sentences = [sent.text for sent in doc.sents]  # 把文本拆分成句子

        current_chunk = []  # 临时存储当前片段的句子
        current_length = 0  # 当前片段的长度

        for sentence in sentences:
            sentence_length = len(sentence)

            # 如果加上当前句子会超过chunk_size，就把当前片段保存
            if current_length + sentence_length > chunk_size:
                # 拼接当前片段的所有句子
                chunk_text = "".join(current_chunk).strip()
                if chunk_text:  # 确保片段不为空
                    chunks.append({
                        "page_num": page_num,
                        "chunk": chunk_text
                    })

                # 重叠处理：保留最后几个句子，让片段之间有衔接
                # 从后往前找，直到总长度小于chunk_overlap
                overlap = []
                overlap_length = 0
                for sent in reversed(current_chunk):
                    if overlap_length + len(sent) <= chunk_overlap:
                        overlap.insert(0, sent)
                        overlap_length += len(sent)
                    else:
                        break

                current_chunk = overlap  # 用重叠部分作为新片段的开始
                current_length = overlap_length

            # 添加当前句子到片段
            current_chunk.append(sentence)
            current_length += sentence_length

        # 处理最后一个片段
        if current_chunk:
            chunk_text = "".join(current_chunk).strip()
            if chunk_text:
                chunks.append({
                    "page_num": page_num,
                    "chunk": chunk_text
                })

    return chunks


# 测试代码
if __name__ == "__main__":
    # 先调用上一步的函数读取PDF（这里简化处理，实际使用时需要导入）
    # 假设已经读取了PDF内容
    sample_pdf_contents = [
        {
            "page_num": 1,
            "text": "平安保险成立于1988年，总部位于深圳。旗下有多个保险产品，包括寿险、车险、健康险等。平安福是公司的明星寿险产品，2024年推出了新版。新版平安福增加了重疾保障种类，从120种扩展到150种。同时，轻症赔付次数也有所增加。"
        }
    ]

    # 调用切分函数
    chunks = split_text(sample_pdf_contents, chunk_size=100, chunk_overlap=20)

    # 打印结果
    print(f"切分后得到{len(chunks)}个片段：")
    for i, chunk in enumerate(chunks, 1):
        print(f"\n片段{i}（第{chunk['page_num']}页）：")
        print(chunk['chunk'])
